package geppetto.cat.corpus;

import gnu.trove.TObjectIntHashMap;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;

public class BilingualCorpusWithPOS extends BilingualCorpus {

	public ArrayList<int[]> _trainSourcePOSSentences = new ArrayList<int[]>();
	public ArrayList<int[]> _trainForeignPOSSentences = new ArrayList<int[]>();
	public ArrayList<int[]> _devSourcePOSSentences = new ArrayList<int[]>();
	public ArrayList<int[]> _devForeignPOSSentences = new ArrayList<int[]>();
	public ArrayList<int[]> _testSourcePOSSentences = new ArrayList<int[]>();
	public ArrayList<int[]> _testForeignPOSSentences = new ArrayList<int[]>();
	public TObjectIntHashMap<String> _posVocabSource = new TObjectIntHashMap<String>();
	public TObjectIntHashMap<String> _posVocabForeign = new TObjectIntHashMap<String>();
	String sourcePOSFile;
	String foreignPOSFile;

	public BilingualCorpusWithPOS(String sourcePOSFile, String foreignPOSFile) {
		super();
		this.sourcePOSFile = sourcePOSFile;
		this.foreignPOSFile = foreignPOSFile;
	}

	public BilingualCorpusWithPOS(int maxNumberOfSentences, String name,
			String sourcePrefix, String foreignPrefix) {
		super(maxNumberOfSentences, name, sourcePrefix, foreignPrefix);
	}

	public BilingualCorpusWithPOS reverse() {
		BilingualCorpusWithPOS res = new BilingualCorpusWithPOS(foreignPOSFile, sourcePOSFile);
		res._name = _name;
		res._vocabForeign = _vocabSource;
		res._vocabSource = _vocabForeign;
		res._vocabNamesForeign = _vocabNamesSource;
		res._vocabNamesSource = _vocabNamesForeign;
		res._maxTrainingSentences = _maxTrainingSentences;
		res._trainSourceSentences = _trainForeignSentences;
		res._trainForeignSentences = _trainSourceSentences;
		res._testSourceSentences = _testForeignSentences;
		res._testForeignSentences = _testSourceSentences;
		res._devSourceSentences = _devForeignSentences;
		res._devForeignSentences = _devSourceSentences;
		res._maxTrainingSentences = _maxTrainingSentences;
		res._nrTrainingSentences = _nrTrainingSentences;
		res._maxSentenceSize = _maxSentenceSize;
		res._discardedSentences = _discardedSentences;
		if (_gold != null)
			res._gold = _gold.reverse();
		if (_devGold != null)
			res._devGold = _devGold.reverse();
		res._maxForeignLen = _maxSourceLen;
		res._maxSourceLen = _maxForeignLen;
		res._sourceFertilities = _foreignFertilities;
		res._foreignFertilities = _sourceFertilities;
		res._sourceSuffix = _foreignSuffix;
		res._foreignSuffix = _sourceSuffix;
		res._foreignWordsCounts = _sourceWordsCounts;
		res._sourceWordsCounts = _foreignWordsCounts;
		res._foreignVocabSize = _sourceVocabSize;
		res._sourceVocabSize = _foreignVocabSize;

		res._trainForeignPOSSentences = _trainSourcePOSSentences;
		res._trainSourcePOSSentences = _trainForeignPOSSentences;
		res._devForeignPOSSentences = _devSourcePOSSentences;
		res._devSourcePOSSentences = _devForeignPOSSentences;
		res._testForeignPOSSentences = _testSourcePOSSentences;
		res._testSourcePOSSentences = _testForeignPOSSentences;
		res._posVocabForeign = _posVocabSource;
		res._posVocabSource = _posVocabForeign;

		res.ref = new Reference(res);
		return res;
	}
	
	public void addTrainFile(File sourceFile, File foreignFile,
			int maxSentenceSize) throws IOException {
		int[] result = addFile(sourceFile, foreignFile, maxSentenceSize,
				_nrTrainingSentences, _discardedSentences,
				_maxTrainingSentences, _trainSourceSentences,
				_trainForeignSentences, _trainSourcePOSSentences, _trainForeignPOSSentences);
		_nrTrainingSentences = result[0];
		_discardedSentences = result[1];
	}

	public void addTestFile(String sourceName, String foreignName)
			throws IOException {
		@SuppressWarnings("unused")
		int[] result = addFile(new File(sourceName), new File(foreignName),
				Integer.MAX_VALUE, _testSourceSentences.size(), 0,
				Integer.MAX_VALUE, _testSourceSentences, _testForeignSentences
				, _testSourcePOSSentences, _testForeignPOSSentences);
	}

	public void addDevFile(String sourceName, String foreignName)
			throws IOException {
		@SuppressWarnings("unused")
		int[] result = addFile(new File(sourceName), new File(foreignName),
				Integer.MAX_VALUE, _devSourceSentences.size(), 0,
				Integer.MAX_VALUE, _devSourceSentences, _devForeignSentences
				, _devSourcePOSSentences, _devForeignPOSSentences);
	}

	int[] addFile(File sourceFile, File foreignFile, int maxSentenceSize,
			int numberSentences, int discardedSentences, int maxSentences,
			ArrayList<int[]> sourceSentences,
			ArrayList<int[]> foreignSentences, ArrayList<int[]> sourcePOSList,
			ArrayList<int[]> foreignPOSList) throws IOException {

		if (_maxSentenceSize == 0) {
			_maxSentenceSize = maxSentenceSize;
		} else if (maxSentenceSize <= _maxSentenceSize) {
			_maxSentenceSize = maxSentenceSize;
		}
		// System.out.println("Nr Sentences " + numberSentences + "Discarded " +
		// _discardedSentences);
		int[] result = { numberSentences, discardedSentences };
		if (result[0] >= maxSentences)
			return result;
		BufferedReader sourceReader = new BufferedReader(new InputStreamReader(
				new FileInputStream(sourceFile), "UTF8"));
		BufferedReader foreignReader = new BufferedReader(
				new InputStreamReader(new FileInputStream(foreignFile), "UTF8"));

		boolean hasPOS = new File(sourcePOSFile).exists()
				&& new File(foreignPOSFile).exists();
		if (!hasPOS)
			System.err.println("could not find POS files! These will not be used");

		BufferedReader sourcePOSReader = null;
		BufferedReader foreignPOSReader = null;
		if (hasPOS) {
			sourcePOSReader = new BufferedReader(new InputStreamReader(
					new FileInputStream(sourcePOSFile), "UTF8"));
			foreignPOSReader = new BufferedReader(new InputStreamReader(
					new FileInputStream(foreignPOSFile), "UTF8"));
		}

		String sourceSentence = sourceReader.readLine();
		String foreignSentence = foreignReader.readLine();

		while (sourceSentence != null && foreignSentence != null) {
			if (result[0] >= maxSentences)
				return result;
			String[] sourceTokens = sourceSentence.split("\\s+");
			int sourceLen = sourceTokens.length;
			String[] foreignTokens = foreignSentence.split("\\s+");
			int foreignLen = foreignTokens.length;
			if (sourceLen <= maxSentenceSize && foreignLen <= maxSentenceSize) {
				// System.out.println(sourceLen + " " + _maxSourceLen);
				if (sourceLen > _maxSourceLen)
					_maxSourceLen = sourceLen;
				// System.out.println(foreignLen + " " + _maxForeignLen);
				if (foreignLen > _maxForeignLen)
					_maxForeignLen = foreignLen;
				sourceSentences.add(convertStringToIntArray(sourceTokens,
						sourceLen, _vocabNamesSource, _vocabSource,
						_sourceWordsCounts));
				foreignSentences.add(convertStringToIntArray(foreignTokens,
						foreignLen, _vocabNamesForeign, _vocabForeign,
						_foreignWordsCounts));
				if (hasPOS) {
					String[] sourcePOS = sourcePOSReader.readLine().split("\\s+");
					String[] foreignPOS = foreignPOSReader.readLine()
							.split("\\s+");
					sourcePOSList.add(convertPOSStringToInteger(sourcePOS,
							_posVocabSource));
					foreignPOSList.add(convertPOSStringToInteger(foreignPOS,
							_posVocabForeign));
				}
				result[0]++;
			} else {
				result[1]++;
			}
			sourceSentence = sourceReader.readLine();
			foreignSentence = foreignReader.readLine();
		}
		return result;

	}

	protected int[] convertPOSStringToInteger(String[] pos,
			TObjectIntHashMap<String> dict) {
		int[] ret = new int[pos.length];
		for (int i = 0; i < pos.length; i++) {
			ret[i] = convertPOSStringToInteger(pos[i], dict);
		}
		return ret;
	}

	protected int convertPOSStringToInteger(String pos,
			TObjectIntHashMap<String> dict) {
		if (!dict.contains(pos)) {
			dict.put(pos, dict.keys().length);
		}
		return dict.get(pos);
	}
	
	public int[] getForeignSentencePOS(int sentenceNr, byte sentenceSource) {
		switch (sentenceSource) {
		case TRAIN_CORPUS:
			return _trainForeignPOSSentences.get(sentenceNr);
		case TEST_CORPUS:
			return _testForeignPOSSentences.get(sentenceNr);
		case DEV_CORPUS:
			return _devForeignPOSSentences.get(sentenceNr);
		}
		return null;
	}

	public int[] getSourceSentencePOS(int sentenceNr, byte sentenceSource) {
		switch (sentenceSource) {
		case TRAIN_CORPUS:
			return _trainSourcePOSSentences.get(sentenceNr);
		case TEST_CORPUS:
			return _testSourcePOSSentences.get(sentenceNr);
		case DEV_CORPUS:
			return _devSourcePOSSentences.get(sentenceNr);
		}
		return null;
	}
}
